
from __future__ import print_function
from keras import backend as K

from keras import activations

from keras.layers import *

Num_capsule = 2
Dim_capsule = 2
Routings = 3


def squash(x, axis=-1):
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True) + K.epsilon()
    scale = K.sqrt(s_squared_norm) / (0.5 + s_squared_norm)
    return scale * x


# define our own softmax function instead of K.softmax
# because K.softmax can not specify axis.
# 自定义softmax函数，替换K.softmax函数，因为K.softmax函数不能指定轴
def softmax(x, axis=-1):
    ex = K.exp(x - K.max(x, axis=axis, keepdims=True))
    return ex / K.sum(ex, axis=axis, keepdims=True)


# define the margin loss like hinge loss
# 定义利润边缘损失，如hinge损失
def margin_loss(y_true, y_pred):
    lamb, margin = 0.5, 0.1
    return y_true * K.square(K.relu(1 - margin - y_pred)) + lamb * (
            1 - y_true) * K.square(K.relu(y_pred - margin))


class Capsule(Layer):
    """A Capsule Implement with Pure Keras
    基于纯Keras的胶囊（组神经元）实现
    There are two vesions of Capsule.
    以下是2个版本的胶囊（组神经元）
    One is like dense layer (for the fixed-shape input),
    一种类似与全连接层（对于固定形状输入），
    and the other is like timedistributed dense (for various length input).
    一种类似时间分布的全连接层（对于变长输入），
    The input shape of Capsule must be (batch_size,
                                        input_num_capsule,
                                        input_dim_capsule
                                       )
    胶囊（组神经元）输入（数据）形状为：
    and the output shape is (batch_size,
                             num_capsule,
                             dim_capsule
                            )
     输出（数据）形状为：
    Capsule Implement is from https://github.com/bojone/Capsule/
    胶囊（组神经元）实现见：https://github.com/bojone/Capsule/
    Capsule Paper: https://arxiv.org/abs/1710.09829
    胶囊（组神经元）论文：https://arxiv.org/abs/1710.09829
    """

    def __init__(self,
                 num_capsule,
                 dim_capsule,
                 routings=3,
                 share_weights=True,
                 activation='squash',
                 **kwargs):
        super(Capsule, self).__init__(**kwargs)
        self.num_capsule = num_capsule
        self.dim_capsule = dim_capsule
        self.routings = routings
        self.share_weights = share_weights
        if activation == 'squash':
            self.activation = squash
        else:
            self.activation = activations.get(activation)

    def build(self, input_shape):
        input_dim_capsule = input_shape[-1]
        if self.share_weights:
            self.kernel = self.add_weight(
                name='capsule_kernel',
                shape=(1, input_dim_capsule,
                       self.num_capsule * self.dim_capsule),
                initializer='glorot_uniform',
                trainable=True)
        else:
            input_num_capsule = input_shape[-2]
            self.kernel = self.add_weight(
                name='capsule_kernel',
                shape=(input_num_capsule, input_dim_capsule,
                       self.num_capsule * self.dim_capsule),
                initializer='glorot_uniform',
                trainable=True)

    def call(self, inputs):
        """Following the routing algorithm from Hinton's paper,
        根据Hinton论文的路由算法，
        but replace b = b + <u,v> with b = <u,v>.
        但是 用b = <u,v>替换b = b + <u,v>
        This change can improve the feature representation of Capsule.
        这种改变可以改善胶囊（组神经元）的特征表示。
        However, you can replace
        而且，可以把替换
            b = K.batch_dot(outputs, hat_inputs, [2, 3])
        with
        使用
            b += K.batch_dot(outputs, hat_inputs, [2, 3])
        to realize a standard routing.
        实现标准路由。
        """

        if self.share_weights:
            hat_inputs = K.conv1d(inputs, self.kernel)
        else:
            hat_inputs = K.local_conv1d(inputs, self.kernel, [1], [1])

        batch_size = K.shape(inputs)[0]
        input_num_capsule = K.shape(inputs)[1]
        hat_inputs = K.reshape(hat_inputs,
                               (batch_size, input_num_capsule,
                                self.num_capsule, self.dim_capsule))
        hat_inputs = K.permute_dimensions(hat_inputs, (0, 2, 1, 3))

        b = K.zeros_like(hat_inputs[:, :, :, 0])
        for i in range(self.routings):
            c = softmax(b, 1)
            if K.backend() == 'theano':
                o = K.sum(o, axis=1)
            o = self.activation(K.batch_dot(c, hat_inputs, [2, 2]))
            if i < self.routings - 1:
                b = K.batch_dot(o, hat_inputs, [2, 3])
                if K.backend() == 'theano':
                    o = K.sum(o, axis=1)

        return o

    def compute_output_shape(self, input_shape):
        return (None, self.num_capsule, self.dim_capsule)


